package com.djf.springaidemo.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.CompressionQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.Map;

/**
 * @Author 26393
 * @Date 2025/5/5 1:12
 * @Version 1.0
 * @Description:rag增强
 */
@RestController
@RequestMapping("/rag")
public class RagController {

    private final ChatClient chatClient;
    private final VectorStore vectorStore;
    private final MultiQueryExpander multiQueryExpander;
    private final RewriteQueryTransformer rewriteQueryTransformer;
    private final TranslationQueryTransformer translationQueryTransformer;
    private final CompressionQueryTransformer compressionQueryTransformer;
    private final RetrievalAugmentationAdvisor retrievalAdvisor;

    @Autowired
    public RagController(
            ChatClient chatClient,
            VectorStore vectorStore,
            MultiQueryExpander multiQueryExpander,
            RewriteQueryTransformer rewriteQueryTransformer,
            TranslationQueryTransformer translationQueryTransformer,
            CompressionQueryTransformer compressionQueryTransformer,
            RetrievalAugmentationAdvisor retrievalAdvisor
    ) {
        this.chatClient = chatClient;
        this.vectorStore = vectorStore;
        this.multiQueryExpander = multiQueryExpander;
        this.rewriteQueryTransformer = rewriteQueryTransformer;
        this.translationQueryTransformer = translationQueryTransformer;
        this.compressionQueryTransformer = compressionQueryTransformer;
        this.retrievalAdvisor = retrievalAdvisor;
    }

    // ====================== 基础 RAG 功能 ======================

    /**
     * 基础问答接口
     */
    @PostMapping("/chat")
    public String basicChat(@RequestParam String userInput) {
        return chatClient.prompt()
                .user(userInput)
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .call()
                .content();
    }

    // ====================== 模块化增强功能 ======================

    /**
     * 多查询扩展
     */
    @PostMapping("/query-expansion")
    public List<Query> multiQueryExpansion(@RequestParam String originalQuery) {
        return multiQueryExpander.expand(new Query(originalQuery));
    }

    /**
     * 查询重写
     */
    @PostMapping("/query-rewrite")
    public String queryRewrite(@RequestParam String originalQuery) {
        Query query = new Query(originalQuery);
        Query transformed = rewriteQueryTransformer.transform(query);
        return transformed.text();
    }

    /**
     * 查询翻译（英文转中文示例）
     */
    @PostMapping("/query-translate")
    public String queryTranslation(@RequestParam String englishQuery) {
        Query query = new Query(englishQuery);
        Query translated = translationQueryTransformer.transform(query);
        return translated.text();
    }

    /**
     * 上下文感知查询处理
     */
    @PostMapping("/context-aware-query")
    public String contextAwareQuery(
            @RequestParam String currentQuery,
            @RequestParam String historyUserMessage,
            @RequestParam String historyAssistantReply
    ) {
        Query query = Query.builder()
                .text(currentQuery)
                .history(
                        new UserMessage(historyUserMessage),
                        new AssistantMessage(historyAssistantReply)
                )
                .build();

        Query transformed = compressionQueryTransformer.transform(query);
        return chatClient.prompt()
                .user(transformed.text())
                .advisors(retrievalAdvisor)
                .call()
                .content();
    }

    // ====================== 高级检索与文档处理 ======================


    // ====================== 错误处理与边界情况 ======================

    /**
     * 带友好错误处理的问答接口
     */
    @PostMapping("/safe-chat")
    public String safeChat(@RequestParam String userInput) {
        try {
            return chatClient.prompt()
                    .user(userInput)
                    .advisors(retrievalAdvisor)
                    .call()
                    .content();
        } catch (Exception e) {
            return "很抱歉，当前无法回答您的问题。请提供更多背景信息或调整提问方式。";
        }
    }
}

